package com.gitee.zycra.union.api.common.util;

import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.crypto.BlockCipher;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.io.CipherInputStream;
import org.bouncycastle.crypto.modes.AEADBlockCipher;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.params.AEADParameters;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * @author zycra
 */
@Slf4j
public class SymmetricEncryptUtil extends BCUtil {

    protected static final String SM4_LABEL = "SM4";
    protected static final String AES_LABEL = "AES";
    private static final int BUFFER_SIZE = 4096;
    private static final int MAC_SIZE = 128;
    private static final int IV_SIZE = 12;
    private static final int SM4_KEY_SIZE = 128;
    private static final int AES_KEY_SIZE = 256;
    private static final byte[] IV = new byte[IV_SIZE];

    private static final Map<String, Integer> KEY_SIZE_MAP = new HashMap<>();
    private static final Map<String, BlockCipher> BLOCK_CIPHER_MAP = new HashMap<>();

    static {
        KEY_SIZE_MAP.put(SM4_LABEL, SM4_KEY_SIZE);
        BLOCK_CIPHER_MAP.put(SM4_LABEL, new SM4Engine());
        KEY_SIZE_MAP.put(AES_LABEL, AES_KEY_SIZE);
        BLOCK_CIPHER_MAP.put(AES_LABEL, new AESEngine());
    }

    protected static String encrypt(String content, String password, String label) {
        return doEncrypt(content, password, true, label);
    }

    protected static String decrypt(String content, String password, String label) {
        return doEncrypt(content, password, false, label);
    }

    protected static boolean encrypt(File originFile, File encryptFile, String password, String label) {
        return doEncryptFile(originFile, encryptFile, password, true, label);
    }

    protected static boolean decrypt(File originFile, File decryptFile, String password, String label) {
        return doEncryptFile(originFile, decryptFile, password, false, label);
    }

    protected static String generateKey(String label) {
        try {
            KeyGenerator keygen = KeyGenerator.getInstance(label, BouncyCastleProvider.PROVIDER_NAME);
            keygen.init(KEY_SIZE_MAP.get(label));
            SecretKey originalKey = keygen.generateKey();
            return OneWayEncryptUtil.bytesToHex(originalKey.getEncoded());
        } catch (Exception e) {
            log.error("generate key error", e);
        }
        return null;
    }

    private static String doEncrypt(String content, String password, boolean forEncryption, String label) {
        if (content == null || password == null) {
            return null;
        }
        AEADBlockCipher cipher = initCipher(password, forEncryption, label);
        if (cipher == null) {
            return null;
        }
        byte[] data = forEncryption ? content.getBytes(StandardCharsets.UTF_8) : Base64.getMimeDecoder().decode(content);
        try (ByteArrayInputStream inputStream = new ByteArrayInputStream(data);
             CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
             ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
            byte[] cache = new byte[BUFFER_SIZE];
            int r;
            while ((r = cipherInputStream.read(cache)) != -1) {
                outputStream.write(cache, 0, r);
                outputStream.flush();
            }
            byte[] result = outputStream.toByteArray();
            return forEncryption ? Base64.getEncoder().encodeToString(result) : new String(result, StandardCharsets.UTF_8);
        } catch (Exception e) {
            log.error("do encrypt error", e);
        }
        return null;
    }

    private static boolean doEncryptFile(File originFile, File targetFile, String password, boolean forEncryption, String label) {
        if (originFile == null || !Files.isReadable(originFile.toPath()) || password == null || label == null) {
            log.error("do encrypt file param check failed");
            return false;
        }
        AEADBlockCipher cipher = initCipher(password, forEncryption, label);
        if (cipher == null) {
            return false;
        }
        try (InputStream inputStream = new FileInputStream(originFile);
             CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
             OutputStream outputStream = new FileOutputStream(targetFile)) {
            byte[] cache = new byte[BUFFER_SIZE];
            int r;
            while ((r = cipherInputStream.read(cache)) != -1) {
                outputStream.write(cache, 0, r);
                outputStream.flush();
            }
            return true;
        } catch (Exception e) {
            log.error("do encrypt file error", e);
        }
        return false;
    }

    private static AEADBlockCipher initCipher(String password, boolean forEncryption, String label) {
        try {
            AEADBlockCipher cipher = new GCMBlockCipher(BLOCK_CIPHER_MAP.get(label));
            cipher.init(forEncryption, new AEADParameters(new KeyParameter(OneWayEncryptUtil.hexToBytes(password)), MAC_SIZE, IV));
            return cipher;
        } catch (Exception e) {
            log.error("init cipher error", e);
        }
        return null;
    }
}
